import pickle

import  pandas as pd
import numpy as np
from  sklearn.model_selection import train_test_split
import time
from sklearn.naive_bayes import GaussianNB, BernoulliNB, MultinomialNB

data = pd.read_csv("data/titanic_train.csv")
data["Sex_Cleaned"] = np.where(data["Sex"] == "male", 0, 1)
data["Embarked_Cleaned"] = np.where(data["Embarked"] == "S", 0,
                                    np.where(data["Embarked"] == "C", 1,
                                             np.where(data["Embarked"] == "Q", 2, 3
                                             )))

data = data[[
    "Survived",
    "Pclass",
    "Sex_Cleaned",
    "Age",
    "SibSp",
    "Parch",
    "Fare",
    "Embarked_Cleaned"
]].dropna(axis=0, how="any")


X_train, X_test = train_test_split(data, test_size=0.33, random_state=int(time.time()))



nb = GaussianNB()

features = [
    "Pclass",
    "Sex_Cleaned",
    "Age",
    "SibSp",
    "Parch",
    "Fare",
    "Embarked_Cleaned"
]


Y_train = X_train["Survived"].values
X_train = X_train[features].values

Y_test = X_test["Survived"].values
X_test = X_test[features].values

print(X_train)
print(Y_train)

nb.fit(X_train, Y_train)


res = nb.predict(X_test)
loss = 0
for i in range(0, len(res)):
    if res[i] != Y_test[i]:
        loss = loss + 1
print("用训练集初次训练后，模型识别的准确率为：", 1 - loss/len(res))

for i in range(0, len(X_test)):
    X_train = np.row_stack((X_train, X_test[i]))
    Y_train = np.append(Y_train, Y_test[i])
    nb.fit(X_train, Y_train)
    res = nb.predict(X_test)
    loss = 0
    for j in range(0, len(res)):
        if res[j] != Y_test[j]:
            loss = loss + 1
    print("第", i, "次用测试集矫正训练后，模型识别的准确率为：", 1 - loss/len(res))

s = pickle.dumps(nb)
f = open('models/NaiveBayes.model', "wb+")
f.write(s)
f.close()
print("模型保存完成\n")